import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

NUM_WORDS = 10000
(train_data,train_labels),(test_data,test_labels) = tf.keras.datasets.imdb.load_data(num_words=NUM_WORDS)

def multi_hot_sequences(sequences,dimension):
    # 创建一个全零的形状矩阵 (len(sequences), dimension)
    results = np.zeros((len(sequences),dimension))
    for i ,word_indices in enumerate(sequences):
        results[i,word_indices] = 1.0
    return results

train_data = multi_hot_sequences(train_data,dimension=NUM_WORDS)
test_data = multi_hot_sequences(test_data,dimension=NUM_WORDS)

baseline_model = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(16,activation='relu',input_shape=[NUM_WORDS,]),
        tf.keras.layers.Dense(16,activation='relu'),
        tf.keras.layers.Dense(1,activation='sigmoid')
    ]
)
baseline_model.compile(optimizer='adam',
                       loss='binary_crossentropy',
                       metrics=['accuracy','binary_crossentropy'])
baseline_model.summary()
baseline_history = baseline_model.fit(train_data,
                                      train_labels,
                                      epochs=20,
                                      batch_size=512,
                                      validation_data=(test_data,test_labels),
                                      verbose=2)


smaller_model = tf.keras.Sequential([
    tf.keras.layers.Dense(4,activation='relu',input_shape=(NUM_WORDS,)),
    tf.keras.layers.Dense(4,activation='relu'),
    tf.keras.layers.Dense(1,activation='sigmoid')
])
smaller_model.compile(optimizer='adam',
                      loss='binary_crossentropy',
                      metrics=['accuracy', 'binary_crossentropy'])

smaller_model.summary()
smaller_history = smaller_model.fit(train_data,
                                    train_labels,
                                    epochs=20,
                                    batch_size=512,
                                    validation_data=(test_data, test_labels),
                                    verbose=2)

bigger_model = keras.models.Sequential([
    keras.layers.Dense(512, activation='relu', input_shape=(NUM_WORDS,)),
    keras.layers.Dense(512, activation='relu'),
    keras.layers.Dense(1, activation='sigmoid')
])

bigger_model.compile(optimizer='adam',
                     loss='binary_crossentropy',
                     metrics=['accuracy','binary_crossentropy'])

bigger_model.summary()
bigger_history = bigger_model.fit(train_data, train_labels,
                                  epochs=20,
                                  batch_size=512,
                                  validation_data=(test_data, test_labels),
                                  verbose=2)

def plot_history(histories,key='binary_crossentropy'):
    plt.figure(figsize=(16,10))
    for name,history in histories:
        val = plt.plot(history.epoch,history.history['val_' + key],"--",label=name.title()+"val")
        plt.plot(history.epoch, history.history[key], color=val[0].get_color(),
             label=name.title()+' Train')
    plt.xlabel('Epochs')
    plt.ylabel(key.replace('_',' ').title())
    plt.legend()

    plt.xlim([0,max(history.epoch)])

plot_history([('baseline', baseline_history),
              ('smaller', smaller_history),
              ('bigger', bigger_history)])

